import pandas as pd
from  datetime import datetime,date,timedelta
from pandas.core.frame import DataFrame
import numpy as np
from util.query_data.processed_query_data import ProcessedQueryData
from empyrical import alpha_beta, sharpe_ratio, max_drawdown, annual_return, tail_ratio
import math
#import threading
# import multiprocessing
#from pathos.pools import ProcessPool
import time

class HwUtil:
    def __init__(self):
        #除权除息数据
        self.xr_xd_list_a = None
        #行情数据
        self.trade_date_list_md = None
        #股票代码数据
        self.trade_date_list_info = None
        #行情数据-除权除息后
        self.trade_date_list_md_xd = None
        #交易日数据
        self.md_trade_date = None
        #回测交易日数组
        self.trade_date_list = None
        self.res_info_data = pd.DataFrame(columns=['date','code'])
        self.ret_data = pd.DataFrame(columns=['date','code','buy','sell','return','sum_return'])
        self.trade_stk_hold = None
        self.trade_mtss_day = None
        self.financial_market_capitalization = None
        self.pqd = ProcessedQueryData('jk')
        self.for_day = 3
        self.size_sec = 10
        self.stock_rps = None
        self.md_and_rps_data = None
        
    #获取前N天的交易日
    def get_previous_num_trade_date(self,trade_date_frame:DataFrame,day_date:date,day_num:int):
        # return self.pqd.get_last_n_trade_date(day_date,day_num)
        return trade_date_frame[trade_date_frame.trade_date<day_date].iloc[0-day_num].trade_date

    #获取后N天的交易日
    def get_next_num_trade_date(self,trade_date_frame:DataFrame,day_date:date,day_num:int):
        return trade_date_frame[trade_date_frame.trade_date<day_date].iloc[day_num-1].trade_date
        # return self.pqd.get_next_n_trade_date(day_date,day_num)

    #获取前1天的交易日
    def get_previous_trade_date(self,trade_date_frame:DataFrame,day_date:date):
        return self.get_previous_num_trade_date(trade_date_frame,day_date,1)
        # return self.pqd.get_last_n_trade_date(day_date,1)

    #获取后1天的交易日
    def get_next_trade_date(self,trade_date_frame:DataFrame,day_date:date):
        return self.get_next_num_trade_date(trade_date_frame,day_date,1)
        # return self.pqd.get_next_n_trade_date(day_date,1)

    #初始化数据
    def init_md_trade_sec_data(self,sta_date_str:str,end_date_str:str,for_day:int,is_fmc:bool,is_stk:bool,is_mtss:bool):
        self.md_trade_date = self.pqd.get_trade_date('2010-01-01','2025-01-01')
        sta_date = datetime.strptime(sta_date_str, "%Y-%m-%d").date()
        sta_date_a = self.get_previous_num_trade_date(self.md_trade_date,sta_date,for_day)

        print("查询行情")
        self.trade_date_list_md = self.pqd.get_md_bar_day_by_trade_date(sta_date_a,end_date_str)
        # self.trade_date_list_md = self.trade_date_list_md.drop_duplicates(subset=['trade_date','code'], keep='first')

        self.trade_date_list_info = self.pqd.get_code_info_by_trade_date(sta_date_a,end_date_str)
        self.trade_date_list_info = self.trade_date_list_info.loc[self.trade_date_list_info['name'].str.contains('ST')==False]
        self.trade_date_list_info = self.trade_date_list_info.loc[(self.trade_date_list_info.code.str[0:1]=='3')==False]
        self.trade_date_list_info = self.trade_date_list_info.loc[(self.trade_date_list_info.code.str[0:3]=='688')==False]
        self.trade_date_list_info = self.trade_date_list_info.loc[self.trade_date_list_info['listed_date']<datetime.strptime('2021-04-01', "%Y-%m-%d").date()]
        print("查询行情-完成")

        print("开始交易日")
        self.trade_date_list = self.pqd.get_trade_date(sta_date_str,end_date_str)['trade_date'].tolist()
        print("开始交易日-完成")

        if is_fmc==True:
            print("获取财务数据")
            self.financial_market_capitalization = self.pqd.get_financial_market_capitalization_by_trade_date(sta_date_a,end_date_str)
            print("获取财务数据-完成")

        if is_stk==True:
            print("获取北向资金持仓")
            self.trade_stk_hold = self.pqd.get_stk_hold_info_by_trade_date(sta_date_a,end_date_str)
            print("获取北向资金持仓-完成")
        if is_mtss==True:
            print("获取融资融卷数据")
            self.trade_mtss_day = self.pqd.get_mtss_info_by_trade_date(sta_date_a,end_date_str)
            print("获取融资融卷数据-完成")

        que_dict = {f"trade_date":">='"+str(sta_date_a)+"' and trade_date<='"+end_date_str+"'"}
        stocks_rps_data = self.pqd.ck.read_dataframe("cqzp_quan_stocks",'stock_rps',que_dict)
        stocks_rps_data = stocks_rps_data[['trade_date','code','RPS5','RPS10','RPS20','RPS50','RPS120','RPS250']]
        stocks_rps_data['code'] = stocks_rps_data['code'].str.replace('XSHE', 'SZ')
        stocks_rps_data['code'] = stocks_rps_data['code'].str.replace('XSHG', 'SH')
        self.stock_rps = stocks_rps_data
        self.stock_rps = self.stock_rps.loc[(self.stock_rps.code.str[0:1]=='3')==False]
        self.stock_rps = self.stock_rps.loc[(self.stock_rps.code.str[0:3]=='688')==False]
        

        print("查询除权除息行情")
        self.trade_date_list_md_xd = self.pqd.get_md_bar_day_xr_xd_by_trade_date(sta_date_a,end_date_str)
        print("查询除权除息行情-完成")

        stock_rps_sindex = self.stock_rps
        stock_rps_sindex['trade_date'] = pd.to_datetime(stock_rps_sindex['trade_date'].astype(str))
        stock_rps_sindex = stock_rps_sindex.set_index(["trade_date","code"])
        trade_date_list_md_xd_sindex = self.trade_date_list_md_xd
        trade_date_list_md_xd_sindex['trade_date'] = pd.to_datetime(trade_date_list_md_xd_sindex['trade_date'].astype(str))
        trade_date_list_md_xd_sindex = trade_date_list_md_xd_sindex.set_index(["trade_date","code"])
        self.md_and_rps_data = pd.concat([stock_rps_sindex,trade_date_list_md_xd_sindex],axis=1)
        self.md_and_rps_data = self.md_and_rps_data.reset_index()

    def _fit(self,data_x, data_y):
        """
        linear regression
        :param data_x:
        :param data_y:
        :return:
        """
        m = len(data_y)
        s = 0
        for i in data_x:
            s += i
        x_bar = s / len(data_x)
        sum_yx = 0
        sum_x2 = 0
        sum_delta = 0
        for i in range(m):
            x = data_x[i]
            y = data_y[i]
            sum_yx += y * (x - x_bar)
            sum_x2 += x ** 2
        # 根据公式计算w
        w = sum_yx / (sum_x2 - m * (x_bar ** 2))

        for i in range(m):
            x = data_x[i]
            y = data_y[i]
            sum_delta += (y - w * x)
        b = sum_delta / m
        return w, b


    def fit(self,data_x: pd.Series, data_y: pd.Series):
        # ls_x = List(data_x.tolist())
        # ls_y = List(data_y.tolist())
        ls_x = data_x
        ls_y = data_y
        res = self._fit(ls_x, ls_y)
        return res[0]

    def threads_for(self,select_sec_data,groupby_code,code_list,select_date_str,select_date,for_fun):
        try:
            for index in range(len(code_list)):
                #从股票数据中查询今天是否交易
                key = code_list[index]
                trade_date_list_info_c_key = self.trade_date_list_info[(self.trade_date_list_info['code']==key)&(self.trade_date_list_info['trade_date']==datetime.strptime(select_date, "%Y-%m-%d").date())]
                if len(trade_date_list_info_c_key)<=0:
                    continue
                select_sec_md_data = groupby_code.get_group(key)
                select_sec_md_data = select_sec_md_data.sort_values('trade_date')
                if len(select_sec_md_data)==0:
                    continue
                select_sec_data_apple = for_fun(select_sec_md_data,select_date_str)
                if select_sec_data_apple is not None:
                    select_sec_data = select_sec_data.append(select_sec_data_apple)
        except Exception as e:
            print(e)

    # 选股方法
    # select_date 选股的日期
    # hu 数据存储类
    # for_day 需要从选股日期往前多少个交易日的数据
    #
    #
    # 当这个交易和这个code符合因子选股条件 
    # 就执行添加到select_sec_data这个dataframe
    # select_sec_data.loc[select_sec_data.shape[0]] = {'code':key,'pre_close':sec_data_0_md.close,'sort':compare_df_sum}
    # 最后返回
    def select_sec_list(self,select_date:date,for_day:int,for_fun):
        #创建返回datafrmae
        colums = ['code','pre_close','sort']
        select_sec_data = pd.DataFrame(columns=colums)
        select_date_str = select_date.strftime('%Y-%m-%d')
        #获得往前N天的交易日和前一个交易日 并筛选出未复权的日K数据
        sta_date_que = self.get_previous_num_trade_date(self.md_trade_date,select_date,for_day)
        end_date_que = self.get_previous_trade_date(self.md_trade_date,select_date)
        select_md_data = self.trade_date_list_md[(self.trade_date_list_md['trade_date']>=sta_date_que)&(self.trade_date_list_md['trade_date']<=end_date_que)]
        #分组
        groupby_code = select_md_data.groupby("code")
        code_list = list(groupby_code.groups)

        for index in range(len(code_list)):
            #从股票数据中查询今天是否交易
            key = code_list[index]
            trade_date_list_info_c_key = self.trade_date_list_info[(self.trade_date_list_info['code']==key)&(self.trade_date_list_info['trade_date']==select_date)]
            if len(trade_date_list_info_c_key)<=0:
                continue
            select_sec_md_data = groupby_code.get_group(key)
            select_sec_md_data = select_sec_md_data.sort_values('trade_date')
            if len(select_sec_md_data)==0:
                continue
            select_sec_data_apple = for_fun(select_sec_md_data,select_date_str)
            if select_sec_data_apple is not None:
                select_sec_data = select_sec_data.append(select_sec_data_apple)
        return select_sec_data
        # # 多线程部分
        # a_len = math.floor(len(code_list)/3)
        # code_list_a = code_list[0:a_len]
        # code_list_b = code_list[a_len:a_len*2-1]
        # code_list_c = code_list[a_len*2:len(code_list)-1]
        # threads = []
        # t1 = threading.Thread(target=self.threads_for,args=(select_sec_data,groupby_code,code_list_a,select_date_str,select_date,for_fun))
        # threads.append(t1)
        # t2 = threading.Thread(target=self.threads_for,args=(select_sec_data,groupby_code,code_list_b,select_date_str,select_date,for_fun))
        # threads.append(t2)
        # t3 = threading.Thread(target=self.threads_for,args=(select_sec_data,groupby_code,code_list_c,select_date_str,select_date,for_fun))
        # threads.append(t3)
        # for t in threads:
        #     t.setDaemon(True)
        #     t.start()
        # t.join()
    
    def process_pool_for(self,size_sec,for_day,for_fun,trade_list):
        try:
            res_info_data = pd.DataFrame(columns=['date','code','size'])
            for td in trade_list:
                th_day = td.strftime('%Y-%m-%d')
                print(th_day)
                select_sec_data = self.select_sec_list(td,for_day,for_fun)
                if len(select_sec_data)==0:
                    continue
                select_size = select_sec_data.size
                if size_sec!=0:
                    if len(select_sec_data)>size_sec:
                        select_sec_data.sort_values("sort",inplace=True,ascending=False)
                        # select_size = select_sec_data.size
                        select_sec_data = select_sec_data[0:size_sec]
                res_info_data.loc[res_info_data.shape[0]] = {'date':th_day,'code':select_sec_data['code'].tolist(),'size':select_size}
            return res_info_data
        except Exception as e:
            print(e)
    #循环选股
    def start_test(self,size_sec:int,for_day:int,for_fun):
        del self.res_info_data
        self.res_info_data = pd.DataFrame(columns=['date','code'])
        # a_len = math.floor(len(self.trade_date_list)/4)
        # trade_list_a = self.trade_date_list[0:a_len]
        # trade_list_b = self.trade_date_list[a_len*1:a_len*2-1]
        # trade_list_c = self.trade_date_list[a_len*2:a_len*3-1]
        # trade_list_d = self.trade_date_list[a_len*3:len(self.trade_date_list)-1]
        # pool = ProcessPool(nodes=4)
        # results = pool.amap(self.process_pool_for,
        # [size_sec,size_sec,size_sec,size_sec], 
        # [for_day,for_day,for_day,for_day], 
        # [for_fun,for_fun,for_fun,for_fun],
        # [trade_list_a,trade_list_b,trade_list_c,trade_list_d])
        # while not results.ready():
        #     time.sleep(1);
        # for res_data in results.get():
        #     self.res_info_data = self.res_info_data.append(res_data)
        self.res_info_data = self.process_pool_for(size_sec,for_day,for_fun,self.trade_date_list)
        

    # #收益统计筛选
    # def sec_filter(self,pre_md_data):
    #     return True

    #num_sell_day:几天后卖出
    #test_money:测试今天
    #day_num:几天后卖出
    #sell_par:卖出入字段 open close
    #buy_par:买入字段 open close
    #hold_money_r:持仓占比 0.5 or 1.0
    def start_test_profit(self,num_sell_day:int,test_money:float,day_num:int,sell_par:str,buy_par:str,hold_money_r:float,sec_filter):
        print("开始收益统计")
        del self.ret_data
        buy_fee = 0.0002
        sell_fee = 0.0012
        self.ret_data = pd.DataFrame(columns=['date','code','buy','sell','return','sum_return','buy_fee','sell_fee'])
        g_money = test_money
        for index, row in self.res_info_data.iterrows():
            date = datetime.strptime(row['date'], "%Y-%m-%d").date()
            print(date)
            code_list = row['code']
            p_sec_num_buy = pd.DataFrame(columns=['code','buy_money','sell_money'])
            for s_code in code_list:
                to_day_md_data = self.trade_date_list_md_xd[(self.trade_date_list_md_xd['trade_date']==date)&(self.trade_date_list_md_xd['code']==s_code)]
                next_md_data = self.trade_date_list_md_xd[(self.trade_date_list_md_xd['trade_date']>date)&(self.trade_date_list_md_xd['code']==s_code)]
                pre_md_data = self.trade_date_list_md_xd[(self.trade_date_list_md_xd['trade_date']<date)&(self.trade_date_list_md_xd['code']==s_code)]
                if len(pre_md_data)>day_num:
                    md_num_data = pre_md_data.iloc[0-day_num]
                else:
                    md_num_data = pre_md_data.iloc[0]
                sell_money = 0
                buy_money = 0
                if len(next_md_data)>(num_sell_day-1):
                    next_md_data = next_md_data[num_sell_day-1:num_sell_day]
                    next_md = next_md_data.iloc[0]
                    sell_money = next_md[sell_par]
                if len(to_day_md_data)>0:
                    md = to_day_md_data.iloc[0]
                    buy_money = md[buy_par]
                    # pre = pre_md_data.iloc[-1]
                    # ppre = pre_md_data.iloc[-2]
                    is_sec_filter = sec_filter(pre_md_data,md)
                    if is_sec_filter==True:
                        p_sec_num_buy.loc[p_sec_num_buy.shape[0]] = {'code':md.code,'sell_money':sell_money,'buy_money':buy_money}
                to_day_md_data = None
                next_md_data = None
                sell_money = 0
                buy_money = 0
            p_sec_num_buy['money'] = (g_money*hold_money_r)/len(p_sec_num_buy)
            p_sec_num_buy['vol'] = p_sec_num_buy['money']/(p_sec_num_buy['buy_money']*100)
            p_sec_num_buy['vol'] = round(p_sec_num_buy['vol'],0)*100
            p_sec_num_buy['b_money'] = p_sec_num_buy['buy_money']*p_sec_num_buy['vol']
            p_sec_num_buy['s_money'] = p_sec_num_buy['sell_money']*p_sec_num_buy['vol']
            buy_money = p_sec_num_buy['b_money'].sum()
            sell_money = p_sec_num_buy['s_money'].sum()
            buy_money_fee = buy_money*buy_fee
            sell_money_fee = sell_money*sell_fee
            if buy_money == 0 or sell_money==0:
                self.ret_data.loc[self.ret_data.shape[0]] = {'date':date,'code':code_list,'buy':0,'sell':0,'return':0,"sum_return":g_money,'buy_fee':0,'sell_fee':0}
            else:
                return_money = sell_money-buy_money-buy_money_fee-sell_money_fee
                g_money = g_money+return_money
                self.ret_data.loc[self.ret_data.shape[0]] = {'date':date,'code':code_list,'buy':buy_money,'sell':sell_money,'return':return_money,"sum_return":g_money,'buy_fee':buy_money_fee,'sell_fee':sell_money_fee}
            p_sec_num_buy = None
            buy_money = 0
            sell_money = 0
            return_money = 0
        print("收益统计完成")
        self.p_profit(test_money)

        
    def p_profit(self,test_money:float):
        z_money = test_money
        huret_data = self.ret_data.set_index('date')
        self.ret_data[['date', 'sum_return']].plot(figsize=(16,9), secondary_y='sum_return')
        self.ret_data['returns_a'] = self.ret_data['return']/self.ret_data['sum_return']
        sum_money = self.ret_data.iloc[-1].sum_return
        win = sum_money-z_money
        max_draw = max_drawdown(self.ret_data['returns_a'].values)*100
        ret_r = win/test_money*100
        year_win_r = win/z_money/len(self.ret_data)*365*100
        win_r = len(self.ret_data[self.ret_data['return']>0])/len(self.ret_data)*100
        max_money_use_r = self.ret_data['buy'].max()/self.ret_data[self.ret_data['buy'].max()==self.ret_data['buy']].iloc[0].sum_return*100
        l_max_los = self.ret_data['return'].min()
        l_max_los_r = self.ret_data['returns_a'].min()*100
        l_max_win = self.ret_data['return'].max()
        l_max_win_r = self.ret_data['returns_a'].max()*100
        buy_fee_money = self.ret_data['buy_fee'].sum()
        sell_fee_money = self.ret_data['sell_fee'].sum()
        print("总资金:"+str(round(sum_money,2))+"")
        print("总收益:"+str(round(win,2))+"")
        print("收益率:"+str(round(ret_r,2))+"%")
        print("年化收益率:"+str(round(year_win_r,2))+"%")
        print("最大回撤:"+str(round(max_draw,2))+"%")
        print("胜率:"+str(round(win_r,2))+"%")
        print("最大资金使用率:"+str(round(max_money_use_r,2))+"%")
        print("当次最大亏损:"+str(round(l_max_los,2))+"")
        print("当次最大亏损率:"+str(round(l_max_los_r,2))+"%")
        print("当次最大盈利:"+str(round(l_max_win,2))+"")
        print("当次最大盈利率:"+str(round(l_max_win_r,2))+"%")
        print("总手续费:"+str(round(sell_fee_money+buy_fee_money,2))+"")
        print("买入手续费:"+str(round(buy_fee_money,2))+"")
        print("卖出手续费:"+str(round(sell_fee_money,2))+"")